{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入相关包\n",
    "from utils import DGraphFin\n",
    "from utils.utils import prepare_folder\n",
    "from utils.evaluator import Evaluator\n",
    "import pandas as pd\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "\n",
    "import torch_geometric.transforms as T\n",
    "\n",
    "import numpy as np\n",
    "from torch_geometric.data import Data\n",
    "import os\n",
    "\n",
    "# 随机种子\n",
    "torch.manual_seed(666)\n",
    "\n",
    "#设置gpu设备\n",
    "device = 0\n",
    "device = f'cuda:{device}' if torch.cuda.is_available() else 'cpu'\n",
    "device = torch.device(device)\n",
    "\n",
    "# CUDA is available\n",
    "print(torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义网络模型\n",
    "\n",
    "# 导入需要的包\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch_geometric.transforms as T\n",
    "from torch_geometric.nn import SAGEConv  # 导入 GAT 层\n",
    "\n",
    "# Model Definition\n",
    "class GraphSAGE(torch.nn.Module):\n",
    "    def __init__(self, in_channels, hidden_channels, out_channels, dropout):\n",
    "        super(GraphSAGE, self).__init__()      \n",
    "        self.sageConv1 = SAGEConv(in_channels,hidden_channels)\n",
    "        self.sageConv2 = SAGEConv(hidden_channels,out_channels)\n",
    "        self.dropout   = dropout\n",
    "    \n",
    "    def reset_parameters(self):\n",
    "        self.sageConv1.reset_parameters()\n",
    "        self.sageConv2.reset_parameters()\n",
    "    \n",
    "    def forward(self,x,edge_index):        \n",
    "        x = self.sageConv1(x,edge_index)\n",
    "        x = F.relu(x)\n",
    "        x = F.dropout(x,p=self.dropout,training=self.training)\n",
    "        x = self.sageConv2(x,edge_index)\n",
    "        return F.log_softmax(x,dim=-1)\n",
    "\n",
    "def train(model, data, train_idx, optimizer):\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    out = model(data.x,data.adj_t)[train_idx]\n",
    "    \n",
    "    loss = F.nll_loss(out,data.y[train_idx])\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    return loss.item()\n",
    "\n",
    "def test(model, data, split_idx, evaluator):\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        out = model(data.x, data.adj_t)\n",
    "        y_pred = out.exp()\n",
    "        losses,eval_results = dict(),dict()\n",
    "        for key in['train','valid']:\n",
    "            node_id = split_idx[key]\n",
    "            losses[key] = F.nll_loss(out[node_id], data.y[node_id]).item()\n",
    "            eval_results[key] = evaluator.eval(data.y[node_id], y_pred[node_id])['auc']\n",
    "    return eval_results, losses, y_pred\n",
    "\n",
    "def predict(data,node_id):\n",
    "    \"\"\"\n",
    "    加载模型和模型预测\n",
    "    :param node_id: int, 需要进行预测节点的下标\n",
    "    :return: tensor, 类0以及类1的概率, torch.size[1,2]\n",
    "    \"\"\"\n",
    "    model = GraphSAGE(in_channels = 20, hidden_channels = 128, out_channels = 2, dropout = 0)\n",
    "    model.load_state_dict(torch.load(save_dir+'/model_graphSAGE.pt')) #载入验证集上表现最好的模型\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        out = model(data.x,data.adj_t)[node_id]\n",
    "        y_pred = out.exp()\n",
    "    return y_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path='./datasets/632d74d4e2843a53167ee9a1-momodel/' #数据保存路径\n",
    "save_dir='./results/' #模型保存路径\n",
    "dataset_name='DGraph'\n",
    "dataset = DGraphFin(root=path, name=dataset_name, transform=T.ToSparseTensor())\n",
    "\n",
    "nlabels = dataset.num_classes\n",
    "if dataset_name in ['DGraph']:\n",
    "    nlabels = 2    #本实验中仅需预测类0和类1\n",
    "\n",
    "data = dataset[0]\n",
    "data.adj_t = data.adj_t.to_symmetric() #将有向图转化为无向图\n",
    "\n",
    "if dataset_name in ['DGraph']:\n",
    "    x = data.x\n",
    "    x = (x - x.mean(0)) / x.std(0)\n",
    "    data.x = x\n",
    "if data.y.dim() == 2:\n",
    "    data.y = data.y.squeeze(1)\n",
    "\n",
    "split_idx = {'train': data.train_mask, 'valid': data.valid_mask, 'test': data.test_mask}  #划分训练集，验证集\n",
    "\n",
    "train_idx = split_idx['train']\n",
    "result_dir = prepare_folder(dataset_name,'GraphSAGE')\n",
    "\n",
    "# 查看数据维度\n",
    "print(data)\n",
    "print(data.x.shape)  #feature\n",
    "print(data.y.shape)  #label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义网络模型\n",
    "model = GraphSAGE(in_channels = data.x.size(-1), hidden_channels = 128, out_channels = nlabels, dropout = 0)\n",
    "print('Model GraphSAGE initialized')\n",
    "eval_metric = 'auc'\n",
    "evaluator = Evaluator(eval_metric)\n",
    "epochs = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 训练网络模型\n",
    "import gc\n",
    "gc.collect()\n",
    "print(sum(p.numel() for p in model.parameters()))\n",
    "\n",
    "model.reset_parameters()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-7)\n",
    "min_valid_loss = 1e8\n",
    "\n",
    "Epochs = []\n",
    "Loss   = []\n",
    "Train_AUC = []\n",
    "Valid_AUC = []\n",
    "\n",
    "for epoch in range(1,epochs + 1):\n",
    "    loss = train(model, data, train_idx, optimizer)\n",
    "    eval_results, losses, out = test(model, data, split_idx, evaluator)\n",
    "    train_eval, valid_eval = eval_results['train'], eval_results['valid']\n",
    "    train_loss, valid_loss = losses['train'], losses['valid']\n",
    "\n",
    "    if valid_loss < min_valid_loss:\n",
    "        min_valid_loss = valid_loss\n",
    "        torch.save(model.state_dict(), save_dir+'/model_graphSAGE.pt') #将表现最好的模型保存\n",
    "    \n",
    "    print(f'Epoch: {epoch:02d}, '\n",
    "              f'Loss: {loss:.4f}, '\n",
    "              f'Train: {100 * train_eval:.3f}%, ' # 我们将AUC值乘上100，使其在0-100的区间内\n",
    "              f'Valid: {100 * valid_eval:.3f}% ')\n",
    "\n",
    "    print('{{\"metric\": \"Loss\", \"value\": {:.4f}, \"epoch\": {} }}'.format(loss,epoch))\n",
    "    print('{{\"metric\": \"Train AUC\", \"value\": {:.4f}, \"epoch\": {} }}'.format(100*train_eval,epoch))\n",
    "    print('{{\"metric\": \"Valid AUC\", \"value\": {:.4f}, \"epoch\": {} }}'.format(100*valid_eval,epoch))\n",
    "\n",
    "    Epochs.append(epoch)\n",
    "    Loss.append(loss)\n",
    "    Train_AUC.append(train_eval)\n",
    "    Valid_AUC.append(valid_eval)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataLogger = pd.DataFrame()\n",
    "dataLogger.insert(0,\"Epochs\", Epochs)\n",
    "dataLogger.insert(1,\"Loss\",Loss)\n",
    "dataLogger.insert(2,\"Train_AUC\",Train_AUC)\n",
    "dataLogger.insert(2,\"Valid_AUC\",Valid_AUC)\n",
    "dataLogger.to_excel('GraphSAGE_trainingLog.xlsx',float_format=\"%.4f\",index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "GNN",
   "language": "python",
   "name": "gnn"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
